import numpy as np
import sys
import pickle as pickle
import scipy.spatial.distance as dist
from ast import literal_eval
import json


def load_features(path):
   try:
      pkl_file = open(path, 'rb')
      features = pickle.load(pkl_file)
      return features
   except:
      print('Must provide a pickle file')
      exit()


def distance_euclidean(a,b):
   distance = np.sqrt(np.sum(np.square(a - b)))
   return distance

def sort_dict(path):
    features = load_features(path)
    dict = {}
    for image, feature in features.items():
        num = int(image[12:17])
        dict[num] = (image, feature)
    dict = sorted(dict.items(), key=lambda obj: obj[0])
    return dict



def get_info(path,threshold):
    dict = sort_dict(path)
    frames_features = []
    section = []
    f = 0
    start = 0
    disPre = 0
    featureTmp = []
    count = 1
    frameNum = 5
    for i, (image, feature) in dict:
        dis = distance_euclidean(f, feature)
        print(i,dis)
        if dis - disPre < threshold:
            if count % (frameNum+1) != 0:
                featureTmp.append(feature)
                count += 1
        else:
            count = 1
            frames_features.append(featureTmp)
            featureTmp = []
            end = i
            section.append((start, end - 1))
            start = i
        if i == len(dict) - 1:
            frames_features.append(featureTmp)
            end = i
            section.append((start, end))
        disPre = dis
        f = feature
    return section[1:], frames_features[1:]

def get_meanvalue(feature1,feature2,section):
    values = []
    for i in range(len(feature1)):
        for j in range(len(feature2)):
            distance = distance_euclidean(feature1[i],feature2[j])
            values.append(distance)
    sum = np.mean(values)
    return sum

# def merge_dict(dict):
#

def judge_frame(path,threshold):
    t = 8
    section,features = get_info(path,t)
    print(section)
    current = 0
    num = current + 1
    predict = {}
    while current < len(features):
        temp = [section[current]]
        while num < len(features):
            distance = get_meanvalue(features[current], features[num],section[num])
            print(section[current], section[num], distance)
            if distance < threshold:
                temp.append(section[num])
                del features[num]
                del section[num]
            else:
                num += 1
        predict[current]=temp
        current += 1
        num = current + 1
    return predict


def writejson(dict,path):
    jsObj = json.dumps(dict)
    fileObject = open(path, 'w')
    fileObject.write(jsObj)
    fileObject.close()


if __name__ == '__main__':
    path = 'D:/project/frameRec/f6/5MP4/resnet_v1_101_features.pkl'
    path_jason = 'C:/Users/wenjunli/Desktop/json/6MP4.json'
    # predict = judge_frame(path,26)
    # print(predict)
    #dict = predict[11]
    #writejson(dict,path_jason)
    dict = sort_dict(path)
    f = 0
    for i,(image,feature) in dict:
        if  i == 872:
            f = feature
        if i == 854:
            print(distance_euclidean(feature,f))
        #print(distance_euclidean(features[872],features[854]))



